"""
密码加密工具
使用AES-256加密算法确保数据源密码安全存储
"""

import os
import base64
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from typing import Optional
import logging

logger = logging.getLogger(__name__)


class EncryptionError(Exception):
    """加密相关异常"""
    pass


class DataSourceEncryption:
    """数据源密码加密类"""
    
    def __init__(self):
        self._fernet = None
        self._initialize_encryption()
    
    def _initialize_encryption(self):
        """初始化加密器"""
        try:
            # 从环境变量获取加密密钥
            encryption_key = os.getenv('DATASOURCE_ENCRYPTION_KEY')
            
            if not encryption_key:
                # 如果没有设置环境变量，使用默认密钥（仅用于开发环境）
                encryption_key = 'default-datasource-encryption-key-change-in-production'
                logger.warning("使用默认加密密钥，生产环境请设置DATASOURCE_ENCRYPTION_KEY环境变量")
            
            # 生成密钥
            key = self._derive_key(encryption_key)
            self._fernet = Fernet(key)
            
        except Exception as e:
            logger.error(f"初始化加密器失败: {str(e)}")
            raise EncryptionError(f"初始化加密器失败: {str(e)}")
    
    def _derive_key(self, password: str) -> bytes:
        """从密码派生加密密钥"""
        try:
            # 使用固定的盐值（在生产环境中应该使用随机盐值）
            salt = b'datasource_salt_'
            kdf = PBKDF2HMAC(
                algorithm=hashes.SHA256(),
                length=32,
                salt=salt,
                iterations=100000,
            )
            key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
            return key
        except Exception as e:
            logger.error(f"密钥派生失败: {str(e)}")
            raise EncryptionError(f"密钥派生失败: {str(e)}")
    
    def encrypt(self, plaintext: str) -> str:
        """加密明文密码"""
        try:
            if not plaintext:
                raise ValueError("明文不能为空")
            
            if not self._fernet:
                raise EncryptionError("加密器未初始化")
            
            # 加密
            encrypted_bytes = self._fernet.encrypt(plaintext.encode('utf-8'))
            
            # 转换为base64字符串便于存储
            encrypted_str = base64.urlsafe_b64encode(encrypted_bytes).decode('utf-8')
            
            return encrypted_str
            
        except Exception as e:
            logger.error(f"加密失败: {str(e)}")
            raise EncryptionError(f"加密失败: {str(e)}")
    
    def decrypt(self, encrypted_text: str) -> str:
        """解密密码"""
        try:
            if not encrypted_text:
                raise ValueError("加密文本不能为空")
            
            if not self._fernet:
                raise EncryptionError("加密器未初始化")
            
            # 从base64字符串转换回字节
            encrypted_bytes = base64.urlsafe_b64decode(encrypted_text.encode('utf-8'))
            
            # 解密
            decrypted_bytes = self._fernet.decrypt(encrypted_bytes)
            
            # 转换为字符串
            plaintext = decrypted_bytes.decode('utf-8')
            
            return plaintext
            
        except Exception as e:
            logger.error(f"解密失败: {str(e)}")
            raise EncryptionError(f"解密失败: {str(e)}")
    
    def is_encrypted(self, text: str) -> bool:
        """检查文本是否已加密"""
        try:
            if not text:
                return False
            
            # 尝试解密，如果成功说明是加密的
            self.decrypt(text)
            return True
        except:
            return False


# 创建全局加密实例
_encryption_instance = None


def get_encryption_instance() -> DataSourceEncryption:
    """获取加密实例（单例模式）"""
    global _encryption_instance
    if _encryption_instance is None:
        _encryption_instance = DataSourceEncryption()
    return _encryption_instance


def encrypt_password(password: str) -> str:
    """加密密码的便捷函数"""
    try:
        encryption = get_encryption_instance()
        return encryption.encrypt(password)
    except Exception as e:
        logger.error(f"密码加密失败: {str(e)}")
        raise EncryptionError(f"密码加密失败: {str(e)}")


def decrypt_password(encrypted_password: str) -> str:
    """解密密码的便捷函数"""
    try:
        encryption = get_encryption_instance()
        return encryption.decrypt(encrypted_password)
    except Exception as e:
        logger.error(f"密码解密失败: {str(e)}")
        raise EncryptionError(f"密码解密失败: {str(e)}")


def is_password_encrypted(password: str) -> bool:
    """检查密码是否已加密的便捷函数"""
    try:
        encryption = get_encryption_instance()
        return encryption.is_encrypted(password)
    except Exception as e:
        logger.error(f"检查密码加密状态失败: {str(e)}")
        return False


def ensure_password_encrypted(password: str) -> str:
    """确保密码已加密（如果未加密则加密）"""
    try:
        if is_password_encrypted(password):
            return password
        else:
            return encrypt_password(password)
    except Exception as e:
        logger.error(f"确保密码加密失败: {str(e)}")
        raise EncryptionError(f"确保密码加密失败: {str(e)}")


def mask_password(password: str, mask_char: str = '*', visible_chars: int = 0) -> str:
    """脱敏显示密码"""
    try:
        if not password:
            return ''
        
        if len(password) <= visible_chars:
            return mask_char * len(password)
        
        if visible_chars > 0:
            return password[:visible_chars] + mask_char * (len(password) - visible_chars)
        else:
            return mask_char * len(password)
            
    except Exception as e:
        logger.error(f"密码脱敏失败: {str(e)}")
        return mask_char * 8  # 返回默认脱敏字符串


# 测试函数（仅用于开发和测试）
def test_encryption():
    """测试加密功能"""
    try:
        test_password = "test_password_123"
        
        # 测试加密
        encrypted = encrypt_password(test_password)
        print(f"原始密码: {test_password}")
        print(f"加密后: {encrypted}")
        
        # 测试解密
        decrypted = decrypt_password(encrypted)
        print(f"解密后: {decrypted}")
        
        # 验证
        assert test_password == decrypted, "加密解密不匹配"
        
        # 测试加密状态检查
        assert is_password_encrypted(encrypted), "加密状态检查失败"
        assert not is_password_encrypted(test_password), "明文状态检查失败"
        
        # 测试脱敏
        masked = mask_password(test_password)
        print(f"脱敏后: {masked}")
        
        print("✅ 加密功能测试通过")
        
    except Exception as e:
        print(f"❌ 加密功能测试失败: {str(e)}")
        raise


if __name__ == "__main__":
    # 运行测试
    test_encryption()
